package sketch.ounit.fuzz;

import java.io.File;
import java.lang.reflect.Method;
import java.util.Collection;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.eclipse.jdt.core.dom.Expression;
import org.eclipse.jdt.core.dom.MethodDeclaration;
import org.eclipse.jdt.core.dom.MethodInvocation;
import org.eclipse.jdt.core.dom.Modifier;

import randoop.main.GenInputsAbstract;
import randoop.util.Timer;
import sketch.ast.ASTUtils;
import sketch.generator.JUnitTestCreator;
import sketch.main.CodeToolkit;
import sketch.main.SketchTestOptions;
import sketch.ounit.Values;
import sketch.util.Checker;
import sketch.util.Globals;
import sketch.util.SketchLog;
import sketch.util.SketchUtil;

//XXX this method needs revise
//XXX this class is going to be obsolete
public class SketchGuidedRandomTestGenerator {

	/**
	 * Required inputs:
	 * 1. the file path of the sketch
	 * 2. the junit class name (including package), and junit method name
	 *    (note that setUp / tearDown method not supported yet)
	 * 3. time_limit (an integer value, in seconds)
	 * */
	public final String file_path;
	public final String junit_class;
	//public final String junit_method;
	public final int time_random_limit;
	public final int time_fuzz_limit;
	
	//the output dir
	private String tmp = "./tmp"; //for storing compiled junit_class
	//private String output_dir = "./output"; //for storing the final output
	private final String junit_name_append = "_guided";
	private final String junit_name_prefix = "RandoopTest";
	
	private Map<MethodDeclaration, List<Expression>> expr_map;
	private Map<MethodDeclaration, List<MethodInvocation>> invoke_map;
	private Map<MethodDeclaration, List<Expression[]>> arg_expr_map;
	
	public void setOutputDir(String newDir) {
		Checker.checkNull(newDir, "The new dir could not be null.");
		File dirFile = new File(newDir);
		if(dirFile.exists()) {
			Checker.checkTrue(dirFile.isDirectory(), newDir + " is not a directory.");
		} else {
			dirFile.mkdirs();
		}
	}
	
	private int maxTestNum = 200;
	public void setMaxTestNum(int num) {
		this.maxTestNum = num;
	}
	
	private String output_package_name = null;
//	"sketchguided.tinysql.autogenerated"
	public void setOutputPackageName(String package_name) {
		this.output_package_name = package_name;
	}
	
	private boolean use_rr = false;
	public void setUseRR(boolean use) {
		this.use_rr = use;
	}
	
	public static boolean output_randoop_tests = false;
		
	public SketchGuidedRandomTestGenerator(String filePath, String junitClass,  int random_time, int fuzz_time, List<Class<?>> classes) {
		this(filePath, junitClass, random_time, fuzz_time);
		if(SketchTestOptions.verbose) {
		    System.out.println("start initializing the object pool, time: " + ExistingObjectFuzzer.time_limit);
		}
		ExistingObjectFuzzer.generate_object_pool(classes);
		if(SketchTestOptions.verbose) {
		    System.out.println("Finish pool initialization.");
		}
	}
	
	public SketchGuidedRandomTestGenerator(String filePath, String junitClass, /*String junitMethod,*/ int random_time, int fuzz_time) {
		Checker.checkTrue(random_time > 0, "The random_time must > 0, current value: " + random_time);
		Checker.checkTrue(fuzz_time > 0, "The fuzz_time must > 0, current value: " + fuzz_time);
		this.file_path = filePath;
		this.junit_class = junitClass;
		//this.junit_method = junitMethod;
		this.time_random_limit = random_time;
		this.time_fuzz_limit = fuzz_time;
		//clean the states of fuzzing
		this.clean_state();
		//set up the random time
		setUpRandomTime();
	}
	
	private void setUpRandomTime() {
		ExistingObjectFuzzer.setTimeLimit(this.time_random_limit);
	}
	
	private void clean_state() {
		ExistingObjectFuzzer.clear_state();
		NewObjectFuzzer.clear_state();
	}
	
	public void generate_tests() {
		RandomPlacesExtractor extractor = new RandomPlacesExtractor(this.file_path);
		this.expr_map = extractor.getExprMap();
		this.invoke_map = extractor.getInvokeMap();
		this.arg_expr_map = extractor.getArgExprMap();
		
		if(expr_map.isEmpty()) {
			if(SketchTestOptions.verbose) {
			    System.out.println("No " + Values.random_fuzz_full_class + "." + Values.random_fuzz_method
					+ " is found!");
			}
			Checker.checkTrue(invoke_map.isEmpty(), "The invoke map must be empty.");
			return;
		}
		Checker.checkTrue(expr_map.size() == invoke_map.size(), "Their size are not the same!");
		Checker.checkTrue(expr_map.size() == arg_expr_map.size(), "Their size are not the same!");
		
		//create a dir
		String packageName = extractor.getPackageName();
		String compile_dir = tmp + Globals.fileSep + SketchUtil.transClassNameDotToSlash(packageName);
		SketchUtil.createDir(compile_dir);
		//compile the code
		List<File> files = new LinkedList<File>();
		files.add(new File(this.file_path));
		CodeToolkit.compileToDir(files, compile_dir);
		//execute the junit method reflectively
		//System.out.println("expr map size: " + expr_map.size() + ", obj id: " + System.identityHashCode(expr_map));
		Map<MethodDeclaration, List<List<AssembledSequenceObjectMutation<?>>>>  generated_method_sequences = null;
		if(this.use_rr) {
			generated_method_sequences = this.execute_junit_rr();
		} else {
			generated_method_sequences = this.execute_junit();	
		}
		
		//write out generated methods
		writeGeneratedTests(generated_method_sequences,extractor);
		if(output_randoop_tests) {
		    writeRandoopGeneratedTests();
		}
	}
	
	private void writeGeneratedTests(Map<MethodDeclaration, List<List<AssembledSequenceObjectMutation<?>>>> generated_method_sequences,
			RandomPlacesExtractor extractor) {
		//generate concrete junit tests
		if(SketchTestOptions.verbose) {
		    System.out.println("Sketched method num: " + generated_method_sequences.keySet().size());
		}
		for(MethodDeclaration method : generated_method_sequences.keySet()) {
			//the string builder
			StringBuilder sb = new StringBuilder();
			List<Expression> expressions = this.expr_map.get(method);
			List<MethodInvocation> invokes = this.invoke_map.get(method);
			List<Expression[]> arg_expressions = this.arg_expr_map.get(method);
			Checker.checkTrue(expressions.size() == invokes.size(), "The expression size: "
					+ expressions.size() + ", the invoke size: " + invokes.size());
			List<List<AssembledSequenceObjectMutation<?>>> lists = generated_method_sequences.get(method);
			if(SketchTestOptions.verbose) {
			    System.out.println("size of lists: " + lists.size());
			}
			this.checkEqual(lists);
			if(SketchTestOptions.verbose) {
			    System.out.println("List size is valid.");
			}
			
			int generated_num = 0;
			for(List<AssembledSequenceObjectMutation<?>> list : lists) {
			   List<String> sequences_str = new LinkedList<String>();
			   List<String> result_vars = new LinkedList<String>();
			   //do the renaming here
			   int length_of_created_sequence = 0;
			   int arg_name_count = 0;
			   for(int i = 0; i < list.size(); i++) {
//				   System.out.println("method: " + method.getName());
//				   System.out.println("  list size: " + list.size() + ", expressions size: " + expressions.size());
				   Expression[] exprs = arg_expressions.get(i);
				   String[] arg_names = new String[exprs.length];
				   String[] arg_exprs = new String[exprs.length];
				   for(int exp_index = 0; exp_index < exprs.length; exp_index++) {
					   arg_names[exp_index] = "arg_var" + (arg_name_count++);
					   arg_exprs[exp_index] = exprs[exp_index].toString();
				   }
				   
				   AssembledSequenceObjectMutation<?> a_sequence = list.get(i);
				   a_sequence.setOutputParams(length_of_created_sequence, expressions.get(i).toString(), arg_names, arg_exprs);
//				   sequences_str.add(a_sequence.toCodeString(length_of_created_sequence, expressions.get(i).toString()));
				   sequences_str.add(a_sequence.toCodeString());
				   result_vars.add(a_sequence.getResultVar(length_of_created_sequence));
				   length_of_created_sequence += list.get(i).size();
			   }
			   MethodDeclaration copy = ASTUtils.deepClone(method);
			   copy.getName().setIdentifier(method.getName().getIdentifier() + junit_name_append + (generated_num++));
			   RandomFuzzInstantiator instantiator = new RandomFuzzInstantiator(copy, invokes, sequences_str, result_vars);
			   if(invokes.size() == sequences_str.size()) {
			     String method_str = instantiator.getResult();
			     sb.append(method_str);
			     sb.append(Globals.lineSep);
			   } else {
				   SketchLog.logln("Discard this sequence, invokes.size() != sequences_str.size(), invoke size: "
						   + invokes.size() + ", sequence string size: " + sequences_str.size());
			   }
			}
			//
			if(SketchTestOptions.verbose) {
			    System.out.println("Generated num in lists: " + lists.size());
			}
			
			//write JUnit test
			String output_name = junit_name_prefix + method.getName().getIdentifier() + "GenTest";
			JUnitTestCreator.createJUnitTest(sb.toString(), SketchTestOptions.output_dir, output_name,
					//(extractor.getUnit().getPackage() == null? null : extractor.getUnit().getPackage().getName().toString()),
					output_package_name,
					ASTUtils.importDeclarations(extractor.getUnit()));
		}
		//System.out.println(sb.toString());
	}
	
	private void writeRandoopGeneratedTests() {
		if(output_package_name != null) {
		    GenInputsAbstract.junit_package_name = output_package_name;
		}
		GenInputsAbstract.junit_output_dir = SketchTestOptions.output_dir;
		ExistingObjectFuzzer.outputRandoopGeneratedTests();
	}
	
	private void checkEqual(List<List<AssembledSequenceObjectMutation<?>>> lists) {
		for(int i = 0; i < lists.size() - 1; i++) {//List<AssembledSequence<?>> list : lists) {
			Checker.checkTrue(lists.get(i).size() == lists.get(i+ 1).size(), "Size is not the same, " +
					" the: " + i + " is: " + lists.get(i).size() + ", and the i+ 1 is: " + lists.get(i + 1).size());
		}
	}
	
	/**
	 * Here r, r means in a round-robin manner
	 * */
	private Map<MethodDeclaration, List<List<AssembledSequenceObjectMutation<?>>>> execute_junit_rr() {
		//the pool must be already created
		Checker.checkTrue(ExistingObjectFuzzer.hasInitializePool(), "The object pool in object fuzzer must be initialized.");
		
		Class<?> clazz = CodeToolkit.loadClassFromDir(tmp, this.junit_class);
		Class<?> super_class = clazz.getSuperclass();
		Checker.checkTrue(super_class != null && super_class.getName().equals("junit.framework.TestCase"),
				"The super_class is: " + super_class);
		//that is for return
		Map<MethodDeclaration, List<List<AssembledSequenceObjectMutation<?>>>> generated_lists_map
	        = new LinkedHashMap<MethodDeclaration, List<List<AssembledSequenceObjectMutation<?>>>>();
		Map<MethodDeclaration, Set<List<AssembledSequenceObjectMutation<?>>>> unique_pools_map
		   = new LinkedHashMap<MethodDeclaration, Set<List<AssembledSequenceObjectMutation<?>>>>();
		Map<MethodDeclaration, Method> method_map = new LinkedHashMap<MethodDeclaration, Method>();
		
		//create an instance of a given class
		Object junit_test_obj = null;
		try {
			junit_test_obj = clazz.newInstance();
		} catch(Exception e) {
			throw new RuntimeException(e);
		}
		
		//go through each method in the map
		for(MethodDeclaration method : expr_map.keySet()) {
			//create an empty list
			List<List<AssembledSequenceObjectMutation<?>>> generated_lists = new LinkedList<List<AssembledSequenceObjectMutation<?>>>();
			generated_lists_map.put(method, generated_lists);
			Set<List<AssembledSequenceObjectMutation<?>>> unique_pools = new HashSet<List<AssembledSequenceObjectMutation<?>>>();
			unique_pools_map.put(method, unique_pools);
			Method exec_junit_method = null;
			try {
				exec_junit_method = clazz.getMethod(method.getName().toString());
			} catch (Exception e) {
				throw new RuntimeException(e);
			}
			method_map.put(method, exec_junit_method);
		}
		
		Timer t = new Timer();
		t.startTiming();
		int count = 0;
		while(t.getTimeElapsedMillis() < this.time_fuzz_limit*1000 && count < maxTestNum) {
			for(MethodDeclaration method : expr_map.keySet()) {
				List<List<AssembledSequenceObjectMutation<?>>> generated_lists = generated_lists_map.get(method);
				Set<List<AssembledSequenceObjectMutation<?>>> unique_pools = unique_pools_map.get(method);
				Method exec_junit_method = method_map.get(method);
				Checker.checkNull(generated_lists, "Generated list could not be null.");
				Checker.checkNull(unique_pools, "The unique pool list could not be null.");
				Checker.checkNull(exec_junit_method, "The executed junit method could not be null.");
				exec_junit_method.setAccessible(true);
				Checker.checkTrue(exec_junit_method.isAccessible() && exec_junit_method.getReturnType().equals(void.class)
						&& !Modifier.isStatic(exec_junit_method.getModifiers()), "method: " + exec_junit_method + " is not a junit method.");
				ExistingObjectFuzzer.clearCurrentGeneratedSequence();
				try {
				  exec_junit_method.invoke(junit_test_obj);
				} catch (Exception e) {
					continue;
				}
				List<AssembledSequenceObjectMutation<?>> seqs = ExistingObjectFuzzer.getAndClearCurrentGeneratedSequences();
				if(!unique_pools.contains(seqs)) {
				  count++;
				  generated_lists.add(seqs);
				  unique_pools.add(seqs);
				}
			}
		}
		t.stopTiming();
		
		//clear the large map to reclaim memory at the end
		unique_pools_map.clear();
		return generated_lists_map;
	}
	
	/**
	 * This method execute each JUnit sketches repeatedly until the time expires, or reaching
	 * the number of total generated unit tests.
	 * */
	private Map<MethodDeclaration, List<List<AssembledSequenceObjectMutation<?>>>> execute_junit() {
		Class<?> clazz = CodeToolkit.loadClassFromDir(tmp, this.junit_class);
		Class<?> super_class = clazz.getSuperclass();
		Checker.checkTrue(super_class != null && super_class.getName().equals("junit.framework.TestCase"),
				"The super_class is: " + super_class);
		Map<MethodDeclaration, List<List<AssembledSequenceObjectMutation<?>>>> generated_lists_map
	        = new LinkedHashMap<MethodDeclaration, List<List<AssembledSequenceObjectMutation<?>>>>();
		try {
			Object junit_test_obj = clazz.newInstance();
//			System.out.println("expr map size: " + expr_map.size() + ", obj id: " + System.identityHashCode(expr_map));
			for(MethodDeclaration method : expr_map.keySet()) {
				Set<List<AssembledSequenceObjectMutation<?>>> unique_pools = new HashSet<List<AssembledSequenceObjectMutation<?>>>();
				List<List<AssembledSequenceObjectMutation<?>>> generated_lists = new LinkedList<List<AssembledSequenceObjectMutation<?>>>();
				Method exec_junit_method = clazz.getMethod(method.getName().toString());
				Checker.checkNull(exec_junit_method, "The method does not exist: " + method.getName());
				//set the method to be accessible
				exec_junit_method.setAccessible(true);
				Checker.checkTrue(exec_junit_method.isAccessible() && exec_junit_method.getReturnType().equals(void.class)
						&& !Modifier.isStatic(exec_junit_method.getModifiers()), "method: " + exec_junit_method + " is not a junit method.");
				//start to execute
				ExistingObjectFuzzer.clearCurrentGeneratedSequence();
				ExistingObjectFuzzer.setTimeLimit(this.time_random_limit);
				//NewObjectFuzzer.setTimeLimit(this.time_limit);
				//XXX the timing here is very tricky
				Timer t = new Timer();
				if(SketchTestOptions.verbose) {
				    System.out.println("method: " + method.getName() +  ", at: " + System.currentTimeMillis());
				}
				int count = 0;
				while(t.getTimeElapsedMillis() < this.time_fuzz_limit*1000) {
					//invoke the junit method will implicitly call the ExistingObjectFuzzer
					exec_junit_method.invoke(junit_test_obj);
					if(count > maxTestNum) {
						//System.out.println("count: " + count + ", maxTestNum: " + maxTestNum + ", break?");
						break;
					}
					//System.out.print(".");
					if(!t.isRunning()) {
						t.startTiming();
					}
					List<AssembledSequenceObjectMutation<?>> seqs = ExistingObjectFuzzer.getAndClearCurrentGeneratedSequences();
					if(!unique_pools.contains(seqs)) {
					  count++;
					  generated_lists.add(seqs);
					  unique_pools.add(seqs);
					}
				}
				t.stopTiming();
				generated_lists_map.put(method, generated_lists);
			}
			return generated_lists_map;
			//XXX move the loop inside? not return just continue?
		} catch (Exception e) {
			//throw new RuntimeException(e);
			//an error thrown from the remaining code, instead of the fuzzing place
			if(SketchTestOptions.verbose) {
			  System.out.println("size: " + generated_lists_map.size());
			  System.out.println("exception leads the program return: " + e.getCause().getMessage());
			}
			return generated_lists_map;
		}
	}
}